import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense

# 设置 Tushare token
ts.set_token('40a93ca838c790d9e0c1fdc4a804c6454383d5c436f71db1affa3a00')
pro = ts.pro_api()


# 获取股票数据
def get_stock_data(tickers, start_date, end_date):
    data_list = []
    for ticker in tickers:
        df = pro.daily(ts_code=ticker, start_date=start_date, end_date=end_date)
        df['trade_date'] = pd.to_datetime(df['trade_date'])
        df.set_index('trade_date', inplace=True)
        df.sort_index(inplace=True)
        df = df[['close']]
        df.columns = [ticker]
        data_list.append(df)
    return pd.concat(data_list, axis=1)


# 基于机器学习的量化策略：LSTM预测价格走势
def lstm_strategy(data, ticker):
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(data[ticker].values.reshape(-1, 1))

    X, y = [], []
    for i in range(60, len(scaled_data)):
        X.append(scaled_data[i - 60:i, 0])
        y.append(scaled_data[i, 0])
    X, y = np.array(X), np.array(y)
    X = np.reshape(X, (X.shape[0], X.shape[1], 1))

    model = Sequential()
    model.add(LSTM(units=50, return_sequences=True, input_shape=(X.shape[1], 1)))
    model.add(LSTM(units=50))
    model.add(Dense(units=1))
    model.compile(optimizer='adam', loss='mean_squared_error')

    model.fit(X, y, epochs=20, batch_size=32)  # 增加训练轮数

    inputs = scaler.transform(data[ticker].values[-60:].reshape(-1, 1))
    inputs = inputs.reshape(1, 60, 1)
    predicted_price = model.predict(inputs)
    predicted_price = scaler.inverse_transform(predicted_price)

    last_price = data[ticker].iloc[-1]
    if predicted_price > last_price:
        return 1
    else:
        return -1


# 量化交易系统类
class TradingSystem:
    def __init__(self, data):
        self.data = data
        self.portfolio = pd.DataFrame(index=data.index, columns=['Value'])
        self.portfolio['Value'] = 1.0
        self.stop_loss = -0.15  # 止损比例
        self.take_profit = 0.25  # 止盈比例

    def generate_signals(self):
        signals = pd.DataFrame(index=self.data.index, columns=self.data.columns)
        for ticker in self.data.columns:
            # 为每只股票生成LSTM信号
            signal = lstm_strategy(self.data, ticker)
            signals[ticker] = signal
        return signals

    def backtest(self):
        signals = self.generate_signals()
        returns = self.data.pct_change()
        strategy_returns = (signals.shift(1).fillna(0) * returns).sum(axis=1)

        for i in range(1, len(self.portfolio)):
            current_value = self.portfolio['Value'].iloc[i - 1]
            daily_return = strategy_returns.iloc[i]
            new_value = current_value * (1 + daily_return)

            # 应用止损和止盈
            if new_value < current_value * (1 + self.stop_loss):
                new_value = current_value * (1 + self.stop_loss)
            elif new_value > current_value * (1 + self.take_profit):
                new_value = current_value * (1 + self.take_profit)

            self.portfolio.loc[self.portfolio.index[i], 'Value'] = new_value

    def plot_results(self):
        plt.figure(figsize=(10, 6))
        plt.plot(self.portfolio['Value'])
        plt.title('LSTM Strategy Performance')
        plt.xlabel('Date')
        plt.ylabel('Portfolio Value')
        plt.show()


if __name__ == "__main__":
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

    tickers_list = ['000001.SZ', '600000.SH', '601318.SH', '000002.SZ']
    start_date = '20200101'
    end_date = '20220101'

    data = get_stock_data(tickers_list, start_date, end_date)

    ts = TradingSystem(data)

    ts.backtest()

    ts.plot_results()

    final_value = ts.portfolio['Value'].iloc[-1]
    print("最终的投资组合价值：", final_value)

    if final_value < 1.0:
        print("策略在回测期间表现不佳，投资组合价值下降了。")
    else:
        print("策略在回测期间表现良好，投资组合价值上涨了。")
